#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : dwspark调用.py
# @Author: HCY
# @Date  : 2025/5/16
# @Desc  : 讯飞星火知识库服务dwspark调用示例

from dwspark.config import Config
# 加载系统环境变量：SPARKAI_UID、SPARKAI_APP_ID、SPARKAI_API_KEY、SPARKAI_API_SECRET

# 创建配置对象，替换为您的应用信息
config = Config()
# 引入星火知识库模型
from dwspark.models import KnowledgeBase
# 日志
from loguru import logger
import requests
import os
import time

'''
星火知识库示例
'''
# 创建星火知识库实例
model = KnowledgeBase(config)

# 示例1：上传文档
logger.info('----------上传文档----------')
# 准备测试文件路径
test_file_path = os.path.join(os.path.dirname(__file__), 'test_document.txt')

# 如果测试文件不存在，创建一个简单的测试文件
if not os.path.exists(test_file_path):
    with open(test_file_path, 'w', encoding='utf-8') as f:
        f.write("这是一个测试文档，用于演示星火知识库API的使用。\n星火大模型是讯飞推出的大语言模型，具有强大的自然语言处理能力。")
    logger.info(f"已创建测试文件: {test_file_path}")

# 上传文档
file_id = model.upload_document(test_file_path)
logger.info(f"上传文档结果，文件ID: {file_id}")

# 示例2：获取文档详情
if file_id:
    logger.info('----------获取文档详情----------')
    doc_info = model.get_document_info(file_id)
    logger.info(f"文档详情: {doc_info}")

# 示例3：获取文档分块内容
if file_id:
    logger.info('----------获取文档分块内容----------')
    # 等待2秒，确保文档已经处理完成
    time.sleep(2)
    chunks = model.get_document_chunks(file_id)
    if chunks:
        logger.info(f"文档分块数量: {len(chunks)}")
        for i, chunk in enumerate(chunks[:2]):  # 只显示前两个分块
            logger.info(f"分块 {i+1}: {chunk.get('content', '')[:50]}...")

# 示例4：创建知识库
logger.info('----------创建知识库----------')
# 添加时间戳避免重名
repo_name = f"测试知识库_{int(time.time())}"
repo_id = model.create_repository(repo_name, "这是一个测试用的知识库", "测试,知识库")
logger.info(f"创建知识库结果，知识库ID: {repo_id}")

# 示例5：将文档添加到知识库
if repo_id and file_id:
    logger.info('----------添加文档到知识库----------')
    # 注意：这里需要调用知识库的添加文件方法，这是一个示例实现
    # 实际使用时请参考API文档或SDK实现
    url = f"{model.REPO_API_URL}/repo/file/add"
    headers = model._get_headers()
    body = {
        "repoId": repo_id,
        "fileIds": [file_id]
    }
    
    import requests
    response = requests.post(url, json=body, headers=headers)
    result = model._handle_response(response, "添加文档到知识库")
    logger.info(f"添加文档到知识库结果: {result}")

# 示例6：获取知识库列表
logger.info('----------获取知识库列表----------')
# 注意：这里需要调用知识库的列表方法，这是一个示例实现
# 实际使用时请参考API文档或SDK实现
url = f"{model.REPO_API_URL}/repo/list"
headers = model._get_headers()
body = {
    "currentPage": 1,
    "pageSize": 10
}

response = requests.post(url, json=body, headers=headers)
result = model._handle_response(response, "获取知识库列表")
if result and "data" in result:
    repos = result["data"].get("rows", [])  # 修正字段名从list改为rows
    logger.info(f"知识库数量: {len(repos)}")
    for repo in repos[:2]:  # 只显示前两个知识库
        logger.info(f"知识库: {repo.get('repoName')} (ID: {repo.get('repoId')})")

# 示例7：与知识库对话
if repo_id:
    logger.info('----------与知识库对话----------')
    # 定义回调函数
    def on_message(data):
        payload = data.get("payload", {})
        content = payload.get("content", "")
        logger.info(f"收到回复: {content}")
    
    def on_error(error):
        logger.error(f"对话错误: {error}")
    
    def on_close():
        logger.info("对话结束")
    
    # 发起对话
    query = "星火大模型有什么特点？"
    logger.info(f"提问: {query}")
    model.chat(
        query=query,
        repo_id=repo_id,
        on_message=on_message,
        on_error=on_error,
        on_close=on_close
    )
    
    # 注意：由于chat方法是异步的，这里需要等待一段时间让回调函数有机会执行
    import time
    time.sleep(5)

logger.info('示例执行完成。')